use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::mem;

/// A set of good hash table prime numbers for resizing.
const SET_PRIMES: &[usize] = &[
    193, 389, 769, 1543, 3079, 6151, 12289, 24593, 49157, 98317, 196613, 393241, 786433, 1572869,
    3145739, 6291469, 12582917, 25165843, 50331653, 100663319, 201326611, 402653189, 805306457,
    1610612741,
];

/// A trait for custom hash functions.
pub trait HashFunc<T> {
    fn hash(&self, value: &T) -> u64;
}

/// A trait for custom equality functions.
pub trait EqualFunc<T> {
    fn equal(&self, a: &T, b: &T) -> bool;
}

/// A trait for custom free functions.
pub trait FreeFunc<T> {
    fn free(&self, value: T);
}

/// A default implementation of `HashFunc` using Rust's `Hash` trait.
pub struct DefaultHashFunc;

impl<T: Hash> HashFunc<T> for DefaultHashFunc {
    fn hash(&self, value: &T) -> u64 {
        let mut hasher = DefaultHasher::new();
        value.hash(&mut hasher);
        hasher.finish()
    }
}

/// A default implementation of `EqualFunc` using Rust's `PartialEq` trait.
pub struct DefaultEqualFunc;

impl<T: PartialEq> EqualFunc<T> for DefaultEqualFunc {
    fn equal(&self, a: &T, b: &T) -> bool {
        a == b
    }
}

/// A default implementation of `FreeFunc` that does nothing.
pub struct DefaultFreeFunc;

impl<T> FreeFunc<T> for DefaultFreeFunc {
    fn free(&self, _value: T) {}
}

/// A node in the hash set.
struct SetEntry<T> {
    data: T,
    next: Option<Box<SetEntry<T>>>,
}

/// A hash set implementation.
pub struct Set<T, HF = DefaultHashFunc, EF = DefaultEqualFunc, FF = DefaultFreeFunc> {
    table: Vec<Option<Box<SetEntry<T>>>>,
    entries: usize,
    prime_index: usize,
    hash_func: HF,
    equal_func: EF,
    free_func: FF,
}

impl<T> Set<T, DefaultHashFunc, DefaultEqualFunc, DefaultFreeFunc> {
    /// Creates a new empty set.
    pub fn new() -> Self {
        Self::with_funcs(DefaultHashFunc, DefaultEqualFunc, DefaultFreeFunc)
    }
}

impl<T, HF, EF, FF> Set<T, HF, EF, FF>
where
    HF: HashFunc<T>,
    EF: EqualFunc<T>,
    FF: FreeFunc<T>,
{
    /// Creates a new empty set with custom hash, equality, and free functions.
    pub fn with_funcs(hash_func: HF, equal_func: EF, free_func: FF) -> Self {
        let table_size = SET_PRIMES[0];
        let table = vec![None; table_size];
        Self {
            table,
            entries: 0,
            prime_index: 0,
            hash_func,
            equal_func,
            free_func,
        }
    }

    /// Inserts a value into the set.
    pub fn insert(&mut self, data: T) -> bool {
        if self.entries * 3 > self.table.len() * 2 {
            self.enlarge();
        }

        let index = (self.hash_func.hash(&data) % self.table.len() as u64) as usize;

        if let Some(entry) = self.table[index].as_mut() {
            let mut current = entry;
            while let Some(ref mut next) = current.next {
                if self.equal_func.equal(&data, &next.data) {
                    return false;
                }
                current = next;
            }
            current.next = Some(Box::new(SetEntry { data, next: None }));
        } else {
            self.table[index] = Some(Box::new(SetEntry { data, next: None }));
        }

        self.entries += 1;
        true
    }

    /// Removes a value from the set.
    pub fn remove(&mut self, data: &T) -> bool {
        let index = (self.hash_func.hash(data) % self.table.len() as u64) as usize;

        let mut current = &mut self.table[index];

        while let Some(ref mut entry) = current {
            if self.equal_func.equal(data, &entry.data) {
                let entry = current.take().unwrap();
                self.free_func.free(entry.data);
                *current = entry.next;
                self.entries -= 1;
                return true;
            }
            current = &mut entry.next;
        }

        false
    }

    /// Checks if a value is in the set.
    pub fn contains(&self, data: &T) -> bool {
        let index = (self.hash_func.hash(data) % self.table.len() as u64) as usize;

        let mut current = self.table[index].as_ref();

        while let Some(entry) = current {
            if self.equal_func.equal(data, &entry.data) {
                return true;
            }
            current = entry.next.as_ref();
        }

        false
    }

    /// Returns the number of entries in the set.
    pub fn len(&self) -> usize {
        self.entries
    }

    /// Enlarges the hash table to the next prime size.
    fn enlarge(&mut self) {
        let old_table = mem::replace(&mut self.table, vec![None; self.table.len()]);
        self.entries = 0;

        if self.prime_index + 1 < SET_PRIMES.len() {
            self.prime_index += 1;
        }

        let new_size = SET_PRIMES[self.prime_index];
        self.table.resize(new_size, None);

        for mut entry in old_table.into_iter().flatten() {
            while let Some(e) = entry.next.take() {
                self.insert(entry.data);
                entry = e;
            }
            self.insert(entry.data);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_insert_and_contains() {
        let mut set = Set::new();
        assert!(set.insert(1));
        assert!(set.insert(2));
        assert!(set.contains(&1));
        assert!(set.contains(&2));
        assert!(!set.contains(&3));
    }

    #[test]
    fn test_remove() {
        let mut set = Set::new();
        assert!(set.insert(1));
        assert!(set.insert(2));
        assert!(set.remove(&1));
        assert!(!set.contains(&1));
        assert!(set.contains(&2));
    }

    #[test]
    fn test_len() {
        let mut set = Set::new();
        assert_eq!(set.len(), 0);
        assert!(set.insert(1));
        assert_eq!(set.len(), 1);
        assert!(set.insert(2));
        assert_eq!(set.len(), 2);
        assert!(set.remove(&1));
        assert_eq!(set.len(), 1);
    }
}
